package com.witmore.neutron.ai.cognitive.mcp.config;

import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.mcp.SyncMcpToolCallback;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * Description: WitmoreSyncMcpToolCallbackProvider
 * Author: 豆子高
 * Date: 2025/8/31 02:34
 */
@Slf4j
@Component
public class WitmoreSyncMcpToolCallbackProvider {

    @Autowired
    private List<McpSyncClient> mcpClients;

    private Map<String, McpSyncClient> cacheMcp;

    @Value("${spring.ai.mcp.client.sse.connections.WMS.url:localhost:8096}")
    private String mcpServerUrl;

    public ToolCallback[] getToolCallbacks() {
        if (mcpClients == null || mcpClients.isEmpty()) {
            return new ToolCallback[0];
        }
        if (cacheMcp == null) {
            cacheMcp = mcpClients.stream().collect(Collectors.toMap(mcpClient -> mcpClient.getClientInfo().name(), Function.identity()));
        }
        return cacheMcp.values().stream()
                .flatMap(mcpClient -> {
                    McpSchema.ListToolsResult listToolsResult = null;
                    try {
                        if (!mcpClient.isInitialized()) {
                            mcpClient.initialize();
                        }
                        listToolsResult = mcpClient.listTools();
                    } catch (Exception e) {

                        HttpClientSseClientTransport transport = HttpClientSseClientTransport
                                .builder(mcpServerUrl)
                                .sseEndpoint(mcpServerUrl + "/sse")
                                .build();
                        McpClient.SyncSpec syncSpec = McpClient.sync(transport)
                                .clientInfo(mcpClient.getClientInfo()) // 设置客户端信息
                                .requestTimeout(Duration.ofSeconds(5));
                        McpSyncClient newMcpClient = syncSpec.build();
                        cacheMcp.put(mcpClient.getClientInfo().name(), newMcpClient);
                        newMcpClient.initialize(); // 执行初始化握手
                        listToolsResult = newMcpClient.listTools();
                        log.info("MCP Client 重新连接成功，获取到 {} 个工具", listToolsResult.tools().size());
                        return listToolsResult.tools()
                                .stream()
                                .map(tool -> new SyncMcpToolCallback(newMcpClient, tool));
                    }
                    if (listToolsResult == null) {
                        return Stream.of();
                    }
                    return listToolsResult.tools()
                            .stream()
                            .map(tool -> new SyncMcpToolCallback(mcpClient, tool));
                })
                .toArray(ToolCallback[]::new);
    }

}
